import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap
from config import RESULT_DIR
import os
import seaborn as sns
from analysis import (
    INDUSTRY_STATS,
    GENDER_DIST,
    AGE_DIST,
    PROVINCE_DIST,
    SCHOOL_DIST,
    EDUCATION_DIST
)

save_path=RESULT_DIR
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False


def plot_province_dist(colormap):
    top_20 = PROVINCE_DIST.sort_values("人数", ascending=False).head(20)
    plt.figure(figsize=(10, 8))
    gradient = np.linspace(0, 1, len(top_20))
    bar_colors = colormap(gradient)
    bars = plt.barh(
        top_20['省份'],
        top_20['人数'],
        color=bar_colors
    )
    for bar in bars:
        width = bar.get_width()
        plt.text(width,
                bar.get_y() + bar.get_height()/2,
                f'{int(width)}',
                va='center', ha='left',
                fontsize=10)
    plt.title("富豪出生地分布TOP20", fontsize=16)
    plt.xlabel("人数", fontsize=12)
    plt.ylabel("出生地", fontsize=12)
    plt.gca().invert_yaxis()  
    plt.grid(axis='x', linestyle='--', alpha=0.5)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "出生地分布.png"), dpi=300, bbox_inches="tight")


def plot_school_dist(colormap):
    sorted_data = SCHOOL_DIST.sort_values("人数", ascending=False).head(20)
    plt.figure(figsize=(10, 8))
    gradient = np.linspace(0, 1, len(sorted_data))
    bar_colors = colormap(gradient)
    bars = plt.barh(
        sorted_data['毕业院校'],
        sorted_data['人数'],
        color=bar_colors
    )
    for bar in bars:
        width = bar.get_width()
        plt.text(width,
                bar.get_y() + bar.get_height()/2,
                f'{int(width)}',
                va='center', ha='left',
                fontsize=10)
    plt.title("富豪毕业院校TOP20", fontsize=16)
    plt.xlabel("人数", fontsize=12)
    plt.ylabel("毕业院校", fontsize=12)
    plt.gca().invert_yaxis()  # 反转y轴使降序排列
    plt.grid(axis='x', linestyle='--', alpha=0.5)
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "毕业院校分布.png"), dpi=300, bbox_inches="tight")


def plot_gender_dist():
    plt.figure(figsize=(6, 6))
    plt.pie(
        GENDER_DIST["人数"],
        labels=GENDER_DIST["性别"],
        autopct="%1.1f%%",
        startangle=90,
        colors=["lightskyblue", "lightcoral"],
    )
    plt.title("富豪性别分布")
    plt.savefig(os.path.join(save_path,"性别分布.png"), dpi=300, bbox_inches="tight")


def plot_education_dist():
    plt.figure(figsize=(6, 6))
    plt.pie(
        EDUCATION_DIST["人数"],
        labels=EDUCATION_DIST["学历"],
        autopct="%1.1f%%",
        startangle=90,
        colors = ["#66C2A5", "#FC8D62", "#8DA0CB"]  # 青、橙、紫
    )
    plt.title("富豪学历分布")
    plt.savefig(os.path.join(save_path,"学历分布.png"), dpi=300, bbox_inches="tight")


def plot_age_dist():
    plt.figure(figsize=(9, 6.3))
    heatmap_data = AGE_DIST.set_index('年龄')['人数'].to_frame()  # 单列热力图
    sns.heatmap(
        data=heatmap_data,
        annot=True,               # 显示数值
        fmt="d",                 # 整数格式
        cmap="YlOrRd",  
        linewidths=0.5,          # 单元格边框线宽
        cbar=True,               # 显示颜色条
        cbar_kws={'label': '人数', 'shrink': 0.8}  # 颜色条标签
    )
    plt.title("富豪年龄分布", pad=20)
    plt.xlabel("")  # 隐藏默认x轴标签
    plt.ylabel("年龄") 
    plt.yticks(fontsize=10, rotation=0) 
    plt.tight_layout()
    plt.savefig(os.path.join(save_path, "年龄分布.png"), dpi=300, bbox_inches="tight")


def plot_industry_stats(target,colormap):
    top_20 = INDUSTRY_STATS.sort_values(target, ascending=False).head(20)
    plt.figure(figsize=(10, 8))
    gradient = np.linspace(0, 1, len(top_20))
    bar_colors = colormap(gradient)
    bars = plt.barh(top_20['行业'], top_20[target], color=bar_colors)
    for bar in bars:
        width = bar.get_width()  # 获取柱子的宽度（富豪数量）
        plt.text(width,  # x坐标：柱子右侧 + 0.5（留出间距）
                bar.get_y() + bar.get_height()/2,  # y坐标：柱子中心
                f'{round(width, 4)}',  # 显示数量
                va='center', ha='left')  # 垂直居中，左对齐
    plt.title(f'{target}排名前20的行业', fontsize=16)
    plt.xlabel(target, fontsize=12)
    plt.ylabel('行业', fontsize=12)
    plt.gca().invert_yaxis() # 反转y轴
    plt.grid(axis='x', linestyle='--', alpha=0.5) # 添加网格线
    plt.tight_layout()
    plt.savefig(os.path.join(save_path,f"行业统计-{target}.png"), dpi=300, bbox_inches="tight")

plot_gender_dist()
plot_education_dist()
plot_age_dist()
colors = [(0.3, 0.1, 0.5), (0.8, 0.6, 1.0)]  # 深紫→浅紫
cmap = LinearSegmentedColormap.from_list("edu", colors)
plot_school_dist(cmap)
colors = [(0.1, 0.4, 0.2), (0.6, 0.9, 0.6)]  # 深绿→浅绿
cmap = LinearSegmentedColormap.from_list("geo", colors)
plot_province_dist(cmap)
colors = [(0.3, 0.3, 0.3), (0.9, 0.9, 0.9)]  # 深灰→浅灰
cmap = LinearSegmentedColormap.from_list("scale", colors)
plot_industry_stats('行业规模',cmap)
colors = [(0.5, 0.0, 0.0), (1.0, 0.3, 0.3)]  # 深红→亮红
cmap = LinearSegmentedColormap.from_list("monopoly", colors)
plot_industry_stats('垄断程度',cmap)
colors = [(0.1, 0.3, 0.6), (0.6, 0.8, 1.0)]  # 深蓝→浅蓝（您原有方案）
cmap = LinearSegmentedColormap.from_list("stable", colors)
plot_industry_stats('稳定性',cmap)
colors = [(1.0, 0.3, 0.1), (1.0, 0.9, 0.4)]  # 深橙→亮黄
cmap = LinearSegmentedColormap.from_list("growth", colors)
plot_industry_stats('发展潜力',cmap)
colors = [(0.4, 0.1, 0.6), (0.8, 0.6, 1.0)]  # 深紫→浅紫（与院校区分明度）
cmap = LinearSegmentedColormap.from_list("composite", colors)
plot_industry_stats('综合发展态势',cmap)